import numpy as np
from scipy.spatial.transform import Rotation

# 读取csv文件并保存为arry数组
def read_csv2arry(filename):
    data_list = []
    with open(filename, encoding='utf-8') as file_obj:
        for line in file_obj:
            data_list.append(line.rstrip().split(','))
    arry = np.array(data_list)
    return arry
    

# 提取数组相关的的数据
def extract_data_2_TUM(arry):
    arry = arry[1:, :]  # 删除第一行  
    extracted_columns_xyz = arry[:, 2:5]    # 提取第3,4,5列xyz数据（索引为2,3,4）
    extracted_columns_rpy = arry[:, 5:8]    # 提取第6,7,8列rpy数据（索引为5,6,7）
    extracted_columns_time = arry[:, 1:2]   # 提取第2列的time数据

    quaternion_array = rpy_array_to_quaternion_array(extracted_columns_rpy)     # 将rpy转换为四元数
    xyz_q = np.hstack((extracted_columns_xyz, quaternion_array))       # 将提取的列与全为0的列进行水平拼接
    result = np.hstack((extracted_columns_time, xyz_q))     # 将提取的time列与其它的列进行水平拼接
    result = result.astype(np.float64)
    result[:, 0] += 1684320000.00
    return result

# 提取数据相关的RPY旋转角并转换为四元数数据
def rpy_array_to_quaternion_array(array):
    rotation = Rotation.from_euler('xyz', array)
    quaternion = rotation.as_quat()
    return quaternion

# 将数据保存为TUM格式的文件
def save_tum_fromat(filename, data):
    with open(filename, 'w') as f:
        for row in data:    # 格式化输出
            timestamp, x, y, z, qx, qy, qz, qw = row
            f.write(f"{timestamp} {x} {y} {z} {qx} {qy} {qz} {qw}\n")

# 将RPY数据转换为四元数数据
def rpy_to_quaternion(roll, pitch, yaw):
    """
    将RPY角(弧度)转换为四元数。
    
    参数:
    roll -- 绕X轴的旋转角(弧度)
    pitch -- 绕Y轴的旋转角(弧度)
    yaw -- 绕Z轴的旋转角(弧度)
    
    返回:
    四元数 [x, y, z, w] 的列表
    """
    rotation = Rotation.from_euler('xyz', [roll, pitch, yaw])
    quaternion = rotation.as_quat()
    return quaternion




print("A1gt2TUM is Start")

datafile = '../log/A1/running/gt_0705_running_00.csv'
save_datafile = '../log/A1/running/gt_0705_running_00.txt'


data_arry = read_csv2arry(datafile)
data_result = extract_data_2_TUM(data_arry)
# print(data_result.shape[0])     # 打印行数
# print(data_result.shape[1])     # 打印列数
save_tum_fromat(save_datafile, data_result)

print("A1gt2TUM is Fineshed")
